import argparse
import copy
import os, pickle
import mindspore
import numpy as np
from pathlib import Path
from mind3d.utils.sim_batch_utils import eval_collate, train_collate
from mind3d.models.build_sim_model import build_model, get_config

from mindspore import load_checkpoint, Tensor, export,ops, save_checkpoint
from mindspore import context, nn

from mindspore import dataset as de
from nuscenes.nuscenes import NuScenes
from mind3d.utils.sim_builder import build_dataset


if not os.path.exists('/mind3d/mindir'):
	os.mkdir('/mind3d/mindir')

def export_model(args):
    device_id = int(os.getenv('DEVICE_ID', '1'))
    device_num = int(os.getenv('RANK_SIZE', '1'))
    context.set_context(mode=context.PYNATIVE_MODE, device_target=args.device_target, device_id=device_id)
    cfg_path = Path(args.config)
    cfg = get_config(cfg_path)

    ms_model = build_model(model_cfg=cfg['model'])

    ckpt = args.checkpoint

    dataset = build_dataset(cfg['data']['train'])
    ds = de.GeneratorDataset(dataset, column_names=cfg['train_column_names'], shuffle=False)
    ds = ds.batch(batch_size=1, input_columns=cfg['train_column_names'], drop_remainder=True, per_batch_map=train_collate)
    iterator = ds.create_dict_iterator()
    expand_dims = ops.ExpandDims()
    data_batch=next(iterator)
    data_batch['points'] = data_batch['points'][0,:,:]
    data_batch['voxels'] = data_batch['voxels'][0,:,:,:]
    data_batch['num_points'] = data_batch['num_points'][0,:]
    data_batch['coordinates'] = data_batch['coordinates'][0,:,:]
    data_batch['num_voxels'] = data_batch['num_voxels'].squeeze(axis=1).astype('int32')
    data_batch['shape'] = data_batch['shape'].astype('int32')
    data_batch['anno_box'] = data_batch['anno_box'].transpose((1, 0 ,2, 3))
    data_batch['ind'] = data_batch['ind'].transpose((1, 0 ,2)).astype('int32')
    data_batch['mask'] = data_batch['mask'].transpose((1, 0 ,2))
    data_batch['cat'] = data_batch['cat'].transpose((1, 0 ,2)).astype('int32')
    data_batch['return_loss']=False
    data_batch['return_feature']=True
    loss=ms_model(data_batch)
    export(ms_model, data_batch, Tensor(True), Tensor(False), file_name='/mind3d/word_dirs/train/new2_copy2', file_format=args.file_format)
    print("success")
    
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", help="train config file path",
                        default='/mind3d/configs/simtrack/simtrack.yaml')
    parser.add_argument('--device_target', default='GPU', help='device id')
    parser.add_argument('--file_format', type=str, default='MINDIR', help="export file format") #
    parser.add_argument(
        "--checkpoint", help="the dir to checkpoint which the model read from",
        default='/mind3d/word_dirs/train/new2_copy2/epoch_18.ckpt')
    
    args = parser.parse_args()
    export_model(args)



